import os
import tensorflow as tf

from code.tfrecords import read_from_tfrecords


class TrainData(object):
    """
    训练类，使用前需要手动open,使用完毕需要手动close
    """

    def __init__(self):
        self._tfrecords_path = None
        self._binary_classfication = None
        self._sess = None
        self._coord = None
        self._threads = None
        self._read_op = None

    def open(self, tfrecords_path, shuffle=True, binary_classfication=False,
             batch_size=32):
        """
        开启会话，开始读取数据
        :param tfrecords_path: 待读取的tfrecords的所在文件夹路径
        :param shuffle: 是否乱序
        :param binary_classfication: 是否二分类
        :param batch_size: 每次获取的数量
        :return:None
        """
        self._tfrecords_path = tfrecords_path
        self._binary_classfication = binary_classfication
        self._sess = tf.Session()
        self._read_op = read_from_tfrecords(self._tfrecords_path, batch_size)
        # 创建一个线程协调器，管理线程
        self._coord = tf.train.Coordinator()
        self._threads = tf.train.start_queue_runners(self._sess, coord=self._coord)

    def get_batch(self):
        if self._sess is None:
            raise AttributeError("请先调用TrainData.open()")
        examples, ids = self._sess.run(self._read_op)
        # 如果binary_classfication = true，即二分类，那么1-10到id都应变为1
        if self._binary_classfication:
            ids = [([1] if id[0] >= 1 else [0]) for id in ids]
        return examples, ids

    def close(self):
        # 需要手动关闭会话
        self._coord.request_stop()
        # 等待线程完成
        self._coord.join(self._threads)
        self._sess.close()

    def get_len(self):
        """
        返回batch大小
        :return: int
        """
        files_count = len(os.listdir(self._tfrecords_path))
        return files_count * 8 * 440
